## 对yolo数据集划分数据集
import os
import shutil
import random
random.seed(0)
import numpy as np
from sklearn.model_selection import train_test_split

# 定义基础绝对路径
base_dir = "/home/shuai/Downloads/my_dataset/select_image/paper_image_label/VOCdevkit"  # 修改为你的实际绝对路径

val_size = 0.1
test_size = 0.2
postfix = 'jpg'
imgpath = os.path.join(base_dir, 'JPEGImages')
txtpath = os.path.join(base_dir, 'txt')

os.makedirs(os.path.join(base_dir, 'images/train'), exist_ok=True)
os.makedirs(os.path.join(base_dir, 'images/val'), exist_ok=True)
os.makedirs(os.path.join(base_dir, 'images/test'), exist_ok=True)
os.makedirs(os.path.join(base_dir, 'labels/train'), exist_ok=True)
os.makedirs(os.path.join(base_dir, 'labels/val'), exist_ok=True)
os.makedirs(os.path.join(base_dir, 'labels/test'), exist_ok=True)

listdir = np.array([i for i in os.listdir(txtpath) if 'txt' in i])
random.shuffle(listdir)
train, val, test = listdir[:int(len(listdir) * (1 - val_size - test_size))], listdir[int(len(listdir) * (1 - val_size - test_size)):int(len(listdir) * (1 - test_size))], listdir[int(len(listdir) * (1 - test_size)):]
print(f'train set size:{len(train)} val set size:{len(val)} test set size:{len(test)}')

for i in train:
    shutil.copy(
        os.path.join(imgpath, f"{i[:-4]}.{postfix}"),
        os.path.join(base_dir, f"images/train/{i[:-4]}.{postfix}")
    )
    shutil.copy(
        os.path.join(txtpath, i),
        os.path.join(base_dir, f"labels/train/{i}")
    )

for i in val:
    shutil.copy(
        os.path.join(imgpath, f"{i[:-4]}.{postfix}"),
        os.path.join(base_dir, f"images/val/{i[:-4]}.{postfix}")
    )
    shutil.copy(
        os.path.join(txtpath, i),
        os.path.join(base_dir, f"labels/val/{i}")
    )

for i in test:
    shutil.copy(
        os.path.join(imgpath, f"{i[:-4]}.{postfix}"),
        os.path.join(base_dir, f"images/test/{i[:-4]}.{postfix}")
    )
    shutil.copy(
        os.path.join(txtpath, i),
        os.path.join(base_dir, f"labels/test/{i}")
    )
